package com.baidu.qianfan.service;

import com.baidu.qianfan.entity.StableDiffusionXLResponseBody;
import com.baidu.qianfan.util.Base64ImageDisplay;
import com.google.gson.Gson;
import okhttp3.*;
import org.json.JSONObject;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import javax.swing.*;
import java.io.IOException;
import java.time.Duration;
import java.util.List;

/**
 * Stable-Diffusion-XL图像ai
 * 百度千帆官网：https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Klkqubb9w
 */
@Service
public class StableDiffusionXL {

    @Value("${baidu.qianfan.API_KEY}")
    private String API_KEY;
    @Value("${baidu.qianfan.SECRET_KEY}")
    private String SECRET_KEY;


    // 发送请求调用接口
    public StableDiffusionXLResponseBody extracted(String text) throws Exception {
        String accessToken = getAccessToken();
        String url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/text2image/sd_xl?access_token=" + accessToken;

        JSONObject jsonObject = new JSONObject();
        jsonObject.put("prompt", text);
        jsonObject.put("size", "1024x1024");
        jsonObject.put("n", 3);
        jsonObject.put("steps", 20);
        jsonObject.put("sampler_index", "DPM++ 2M SDE Karras");

        MediaType JSON = MediaType.parse("application/json; charset=utf-8");
        RequestBody body = RequestBody.create(jsonObject.toString(), JSON);

        OkHttpClient client = new OkHttpClient.Builder()
                .connectTimeout(Duration.ofSeconds(120)) // 连接超时时间 秒
                .writeTimeout(Duration.ofSeconds(120))   // 写入超时时间
                .readTimeout(Duration.ofSeconds(120))    // 读取超时时间
                .build();

        Request request = new Request.Builder()
                .url(url)
                .post(body)
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "application/json")
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) throw new IOException("意外代码：" + response);
            // 读取响应体为字符串
            String responseBody = response.body().string();
            // 使用Gson将字符串转换为Java对象
            Gson gson = new Gson();
            StableDiffusionXLResponseBody stableDiffusionXLResponseBody = gson.fromJson(responseBody, StableDiffusionXLResponseBody.class);
            return stableDiffusionXLResponseBody;
        }
    }

    // 获取token
    private String getAccessToken() throws Exception {
        OkHttpClient client = new OkHttpClient();
        MediaType FORM = MediaType.parse("application/x-www-form-urlencoded; charset=utf-8");

        RequestBody formBody = new FormBody.Builder()
                .add("grant_type", "client_credentials")
                .add("client_id", API_KEY)
                .add("client_secret", SECRET_KEY)
                .build();

        Request request = new Request.Builder()
                .url("https://aip.baidubce.com/oauth/2.0/token")
                .post(formBody)
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) throw new IOException("意外代码：" + response);
            // 将响应结果返回
            JSONObject jsonObject = new JSONObject(response.body().string());
            return jsonObject.getString("access_token");
        }
    }

}